import torch
import torch.nn as nn
import numpy as np
import random
from fastapi import FastAPI, HTTPException, APIRouter
from pydantic import BaseModel
from datetime import datetime
import calendar

app = FastAPI()
router = APIRouter()

# CNN+LSTM 预测模型
class CNN_LSTM_Model(nn.Module):
    def __init__(self, input_dim, cnn_out_channels, lstm_hidden_size, output_dim, fixed_length=30):
        super(CNN_LSTM_Model, self).__init__()
        self.fixed_length = fixed_length
        self.conv1d = nn.Conv1d(in_channels=input_dim, out_channels=cnn_out_channels, kernel_size=3, padding=1)
        self.pool = nn.AdaptiveAvgPool1d(output_size=fixed_length)
        self.lstm = nn.LSTM(input_size=cnn_out_channels, hidden_size=lstm_hidden_size, batch_first=True)
        self.fc = nn.Linear(lstm_hidden_size, output_dim)

    def forward(self, x):
        x = [self.conv1d(seq.permute(1, 0).unsqueeze(0)) for seq in x]
        x = [self.pool(seq).squeeze(0).permute(1, 0) for seq in x]
        x = torch.nn.utils.rnn.pad_sequence(x, batch_first=True)
        lstm_out, _ = self.lstm(x)
        output = self.fc(lstm_out[:, -1, :])
        return output

# 训练函数
def train_model(model, data, targets, epochs=5, learning_rate=0.001):
    criterion = nn.MSELoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
    for epoch in range(epochs):
        optimizer.zero_grad()
        outputs = model(data)
        loss = criterion(outputs, torch.tensor(targets, dtype=torch.float32))
        loss.backward()
        optimizer.step()
        print(f"Epoch [{epoch+1}/{epochs}], Loss: {loss.item():.4f}")

# 滚动预测函数
def rolling_forecast(model, initial_data, predict_start_date):
    predict_start_date = datetime.strptime(predict_start_date, "%Y-%m-%d")
    days_in_month = calendar.monthrange(predict_start_date.year, predict_start_date.month)[1]
    remaining_days = days_in_month - predict_start_date.day + 1

    predictions = []
    current_data = torch.tensor(initial_data, dtype=torch.float32)
    for _ in range(remaining_days):
        output = model([current_data]).item()
        predictions.append(output)
        next_step = torch.tensor([[output] * current_data.shape[1]], dtype=torch.float32)
        current_data = torch.cat((current_data, next_step), dim=0)

    return predictions

# 初始化模型
input_dim, cnn_out_channels, lstm_hidden_size, output_dim, fixed_length = 3, 16, 32, 1, 30
model = CNN_LSTM_Model(input_dim, cnn_out_channels, lstm_hidden_size, output_dim, fixed_length)

# FastAPI 请求格式
class TrainRequest(BaseModel):
    data: list
    targets: list

class PredictRequest(BaseModel):
    initial_data: list
    predict_start_date: str

class RequestType(BaseModel):
    mode: str  # "train" or "predict"
    train_data: TrainRequest | None = None
    predict_data: PredictRequest | None = None

# FastAPI 端点
@app.post("/process")
async def process_request(request: RequestType):
    if request.mode == "train":
        if request.train_data is None:
            raise HTTPException(status_code=400, detail="训练模式需要提供训练数据")
        data_tensors = [torch.tensor(seq, dtype=torch.float32) for seq in request.train_data.data]
        target_tensors = [torch.tensor(target, dtype=torch.float32) for target in request.train_data.targets]
        train_model(model, data_tensors, target_tensors)
        return {"message": "训练完成"}

    elif request.mode == "predict":
        if request.predict_data is None:
            raise HTTPException(status_code=400, detail="预测模式需要提供预测数据")
        predictions = rolling_forecast(model, request.predict_data.initial_data, request.predict_data.predict_start_date)
        return {"predictions": predictions}

    else:
        raise HTTPException(status_code=400, detail="mode 必须是 'train' 或 'predict'")


# 将 app 中的路由添加到 router 中
for route in app.routes:
    if hasattr(route, "methods"):
        router.add_api_route(route.path, route.endpoint, methods=route.methods)